/**
    Copyright (c) 2010 yakiimo02
    Distributed under the New BSD License.
    See included license.txt or http://www.yakiimo3d.com/NewBSDLicense.txt
**/

#include "DXUT.h"
#include "DXHelper.h"
#include "ToneMapCS.h"

ToneMapCS        g_toneMapCS;

/**
*/
ToneMapCS::ToneMapCS() : m_pQuadLayout( NULL ), m_pVB( NULL ),
                         m_pQuadVS( NULL ), m_pQuadPS( NULL ), m_pSampleStateLinear( NULL ), m_pTextureSRV( NULL ), m_pConstantBuf( NULL )
{
}

/**
*/
ToneMapCS::~ToneMapCS()
{
}

/**
    @param pd3dDevice    [in]    D3D Device
*/
HRESULT ToneMapCS::OnCreateDevice( ID3D11Device* pd3dDevice )
{
    HRESULT hr = S_OK;
    ID3DBlob* pBlob = NULL;

    // vertex shader
    V_RETURN( CompileShaderFromFile( L"ToneMap/ToneMapCS.hlsl", "QuadVS", "vs_4_0", &pBlob ) );
    V_RETURN( pd3dDevice->CreateVertexShader( pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pQuadVS ) );
    
    // input layout
    const D3D11_INPUT_ELEMENT_DESC quadlayout[] =
    {
        { "POSITION", 0, DXGI_FORMAT_R32G32B32A32_FLOAT, 0, 0, D3D11_INPUT_PER_VERTEX_DATA, 0 },
        { "TEXCOORD", 0, DXGI_FORMAT_R32G32_FLOAT, 0, 16, D3D11_INPUT_PER_VERTEX_DATA, 0 },
    };
    V_RETURN( pd3dDevice->CreateInputLayout( quadlayout, 2, pBlob->GetBufferPointer(), pBlob->GetBufferSize(), &m_pQuadLayout ) );
    SAFE_RELEASE( pBlob );    

    // pixel shader
    V_RETURN( CompileShaderFromFile( L"ToneMap/ToneMapCS.hlsl", "QuadPS", "ps_4_0", &pBlob ) );
    V_RETURN( pd3dDevice->CreatePixelShader( pBlob->GetBufferPointer(), pBlob->GetBufferSize(), NULL, &m_pQuadPS ) );
    SAFE_RELEASE( pBlob );

    // bilinear sampler state
    D3D11_SAMPLER_DESC SamplerDesc;
    ZeroMemory( &SamplerDesc, sizeof(SamplerDesc) );
    SamplerDesc.AddressU = D3D11_TEXTURE_ADDRESS_CLAMP;
    SamplerDesc.AddressV = D3D11_TEXTURE_ADDRESS_CLAMP;
    SamplerDesc.AddressW = D3D11_TEXTURE_ADDRESS_CLAMP;
    SamplerDesc.Filter = D3D11_FILTER_MIN_MAG_MIP_LINEAR;
    V_RETURN( pd3dDevice->CreateSamplerState( &SamplerDesc, &m_pSampleStateLinear ) );

    // full screen quad
    SCREEN_VERTEX svQuad[4];
    svQuad[0].pos = D3DXVECTOR4( -1.0f, 1.0f, 0.5f, 1.0f );
    svQuad[0].tex = D3DXVECTOR2( 0.0f, 0.0f );
    svQuad[1].pos = D3DXVECTOR4( 1.0f, 1.0f, 0.5f, 1.0f );
    svQuad[1].tex = D3DXVECTOR2( 1.0f, 0.0f );
    svQuad[2].pos = D3DXVECTOR4( -1.0f, -1.0f, 0.5f, 1.0f );
    svQuad[2].tex = D3DXVECTOR2( 0.0f, 1.0f );
    svQuad[3].pos = D3DXVECTOR4( 1.0f, -1.0f, 0.5f, 1.0f );
    svQuad[3].tex = D3DXVECTOR2( 1.0f, 1.0f );

    D3D11_BUFFER_DESC vbdesc =
    {
        4 * sizeof( SCREEN_VERTEX ),
        D3D11_USAGE_DEFAULT,
        D3D11_BIND_VERTEX_BUFFER,
        0,
        0
    };
    D3D11_SUBRESOURCE_DATA InitData;
    InitData.pSysMem = svQuad;
    InitData.SysMemPitch = 0;
    InitData.SysMemSlicePitch = 0;
    V_RETURN( pd3dDevice->CreateBuffer( &vbdesc, &InitData, &m_pVB ) );

    // Setup constant buffers
    D3D11_BUFFER_DESC Desc;
    Desc.Usage = D3D11_USAGE_DYNAMIC;
    Desc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
    Desc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
    Desc.MiscFlags = 0;    
    Desc.ByteWidth = sizeof( ConstantBuf );
    V_RETURN( pd3dDevice->CreateBuffer( &Desc, NULL, &m_pConstantBuf ) );

    V_RETURN( m_measureLuminance.OnCreateDevice( pd3dDevice ) );

    return hr;
}

/**
*/
HRESULT ToneMapCS::OnResizedSwapChain( ID3D11Device* pd3dDevice, const int nWidth, const int nHeight )
{
    return m_measureLuminance.OnResizedSwapChain( pd3dDevice, nWidth, nHeight );
}

/**
*/
void ToneMapCS::OnDestroyDevice()
{
    SAFE_RELEASE( m_pQuadPS );
    SAFE_RELEASE( m_pSampleStateLinear );
    SAFE_RELEASE( m_pVB );
    SAFE_RELEASE( m_pQuadVS );
    SAFE_RELEASE( m_pQuadLayout );    
    SAFE_RELEASE( m_pConstantBuf );

    m_measureLuminance.OnDestroyDevice();
}

/**
*/
void ToneMapCS::OnReleasingSwapChain()
{
    m_measureLuminance.OnReleasingSwapChain();
}

/**
    @param pd3dDevice                [in]    D3D Device
    @param pd3dImmediateContext        [in]    D3D Device Immediate Context
    @param pBackBufferDesc            [in]    Description of backbuffer
    @param nWidth                    [in]    input image width
    @param nHeight                    [in]    input image height
*/
HRESULT ToneMapCS::OnFrameRender( ID3D11Device* pd3dDevice, ID3D11DeviceContext* pd3dImmediateContext, const DXGI_SURFACE_DESC* pBackBufferDesc, const int nWidth, const int nHeight )
{
    HRESULT hr = S_OK;

    // calculate log average luminance and max luminance.
    m_measureLuminance.DoMeasureLuminance( pd3dImmediateContext, nWidth, nHeight, m_pTextureSRV );

    m_constantBuf.fLogAvgLuminance = m_measureLuminance.GetLogAvgLum();
    m_constantBuf.fMaxLuminance = m_measureLuminance.GetLumMax();
    m_constantBuf.nInputWidth = nWidth;
    m_constantBuf.nInputHeight = nHeight;

    // set constant buffer
    D3D11_MAPPED_SUBRESOURCE MappedResource;            
    V( pd3dImmediateContext->Map( m_pConstantBuf, 0, D3D11_MAP_WRITE_DISCARD, 0, &MappedResource ) );
    ConstantBuf* p = (ConstantBuf*)MappedResource.pData;
    *p = m_constantBuf;
    pd3dImmediateContext->Unmap( m_pConstantBuf, 0 );
    ID3D11Buffer* ppCB[1] = { m_pConstantBuf };
    pd3dImmediateContext->PSSetConstantBuffers( 0, 1, ppCB );

    DrawFullScreenQuad( pd3dImmediateContext, m_pQuadPS, pBackBufferDesc->Width, pBackBufferDesc->Height );

    return hr;
}

/**
    @param pd3dImmediateContext        [in]    D3D Device Immediate Context
    @param pPS                        [in]    Pixel Shader
    @param nWidth                    [in]    Backbuffer width
    @param nHeight                    [in]    Backbuffer height
*/
void ToneMapCS::DrawFullScreenQuad( ID3D11DeviceContext* pd3dImmediateContext, ID3D11PixelShader* pPS, UINT nWidth, UINT nHeight )
{
    // Save the old viewport
    D3D11_VIEWPORT vpOld[D3D11_VIEWPORT_AND_SCISSORRECT_MAX_INDEX];
    UINT nViewPorts = 1;
    pd3dImmediateContext->RSGetViewports( &nViewPorts, vpOld );

    // Setup the viewport to match the backbuffer
    D3D11_VIEWPORT vp;
    vp.Width = static_cast< float >( nWidth );
    vp.Height = static_cast< float >( nHeight );
    vp.MinDepth = 0.0f;
    vp.MaxDepth = 1.0f;
    vp.TopLeftX = 0;
    vp.TopLeftY = 0;
    pd3dImmediateContext->RSSetViewports( 1, &vp );

    UINT strides = sizeof( SCREEN_VERTEX );
    UINT offsets = 0;
    ID3D11Buffer* pBuffers[1] = { m_pVB };

    pd3dImmediateContext->IASetInputLayout( m_pQuadLayout );
    pd3dImmediateContext->IASetVertexBuffers( 0, 1, pBuffers, &strides, &offsets );
    pd3dImmediateContext->IASetPrimitiveTopology( D3D11_PRIMITIVE_TOPOLOGY_TRIANGLESTRIP );

    ID3D11SamplerState* aSamplers[] = { m_pSampleStateLinear };
    pd3dImmediateContext->PSSetSamplers( 0, 1, aSamplers );

    ID3D11ShaderResourceView* aRViews[ 1 ] = { m_pTextureSRV };
    pd3dImmediateContext->PSSetShaderResources( 0, 1, aRViews );

    pd3dImmediateContext->VSSetShader( m_pQuadVS, NULL, 0 );
    pd3dImmediateContext->PSSetShader( pPS, NULL, 0 );
    pd3dImmediateContext->Draw( 4, 0 );

    // Restore the Old viewport
    pd3dImmediateContext->RSSetViewports( nViewPorts, vpOld );

    //
    ID3D11ShaderResourceView* aRViewsNull[ 1 ] = { NULL };
    pd3dImmediateContext->PSSetShaderResources( 0, 1, aRViewsNull );
}